package ai

import (
	"context"
	"strings"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1"
	"github.com/rilldata/rill/runtime"
)

const DeveloperAgentName = "developer_agent"

type DeveloperAgent struct {
	Runtime *runtime.Runtime
}

var _ Tool[*DeveloperAgentArgs, *DeveloperAgentResult] = (*DeveloperAgent)(nil)

type DeveloperAgentArgs struct {
	Prompt          string `json:"prompt" jsonschema:"User prompt for the developer agent."`
	InitProject     bool   `json:"init_project,omitempty" jsonschema:"Set to true if the prompt is provided as part of project initialization. Should only be used when the project directory contains an empty project (except for the boilerplate generated by UnpackEmpty)."`
	CurrentFilePath string `json:"current_file_path,omitempty" jsonschema:"Optional path to the file that the user is currently viewing/editing."`
}

type DeveloperAgentResult struct {
	Response string `json:"response"`
}

func (t *DeveloperAgent) Spec() *mcp.Tool {
	return &mcp.Tool{
		Name:        DeveloperAgentName,
		Title:       "Developer Agent",
		Description: "Agent that assists with development tasks.",
		Meta: map[string]any{
			"openai/toolInvocation/invoking": "Developing...",
			"openai/toolInvocation/invoked":  "Development completed",
		},
	}
}

func (t *DeveloperAgent) CheckAccess(ctx context.Context) (bool, error) {
	return checkDeveloperAgentAccess(ctx, t.Runtime)
}

func (t *DeveloperAgent) Handler(ctx context.Context, args *DeveloperAgentArgs) (*DeveloperAgentResult, error) {
	// Pre-invoke file listing
	s := GetSession(ctx)
	_, err := s.CallTool(ctx, RoleAssistant, ListFilesName, nil, &ListFilesArgs{})
	if err != nil {
		return nil, err
	}

	// Generate the prompts
	systemPrompt, err := t.systemPrompt(ctx)
	if err != nil {
		return nil, err
	}
	userPrompt, err := t.userPrompt(args)
	if err != nil {
		return nil, err
	}

	// Build initial completion messages
	messages := []*aiv1.CompletionMessage{NewTextCompletionMessage(RoleSystem, systemPrompt)}
	messages = append(messages, s.NewCompletionMessages(s.MessagesWithResults(FilterByRoot()))...)
	messages = append(messages, NewTextCompletionMessage(RoleUser, userPrompt))
	messages = append(messages, s.NewCompletionMessages(s.MessagesWithResults(FilterByParent(s.ID())))...)

	// Run an LLM tool call loop
	var response string
	err = s.Complete(ctx, "Developer loop", &response, &CompleteOptions{
		Messages:      messages,
		Tools:         []string{ListFilesName, SearchFilesName, ReadFileName, DevelopModelName, DevelopMetricsViewName},
		MaxIterations: 10,
		UnwrapCall:    true,
	})
	if err != nil {
		return nil, err
	}

	return &DeveloperAgentResult{
		Response: response,
	}, nil
}

func (t *DeveloperAgent) systemPrompt(ctx context.Context) (string, error) {
	// Prepare template data.
	session := GetSession(ctx)
	data := map[string]any{
		"ai_instructions": session.ProjectInstructions(),
	}

	// Generate the system prompt
	return executeTemplate(`<role>You are a data engineer agent specialized in developing data models and metrics view definitions in the Rill business intelligence platform.</role>

<concepts>
Rill is a "business intelligence as code" platform where all resources are defined using YAML files containing SQL snippets in a project directory.
Rill supports many different resource types, such as connectors, models, metrics views, explore dashboards, canvas dashboards, and more.
For the purposes of your work, you will only deal with:
- **Models**: SQL statements and related metadata that produce a single table in the project's database (usually DuckDB or Clickhouse).
- **Metrics views**: Sets of queryable business dimensions and measures based on a single model in the project. This is sometimes called the "semantic layer" or "metrics layer" in other tools.
Rill maintains a DAG of resources. In this DAG, metrics views are always derived from a single model. Multiple metrics views can derive from a single model, although usually it makes sense to have just one metrics view per model.
When users ask you to develop a "dashboard", that just means to develop a new metrics view (and possibly a new underlying model). Rill automatically creates visual dashboards for each metrics view.
</concepts>

<example>
This example is not directly related to your current task. It just serves to explain how Rill project's look and how you might act on a user request.

Rill projects often (but not always) organize files in directories by resource type. A Rill project might look like this:
{{ backticks }}
connectors/duckdb.yaml
models/orders.yaml
metrics/orders.yaml
rill.yaml
{{ backticks }}

The user might ask you to "Create a dashboard for my Github activity". You would notice that this does not relate to the current files, and proceed with the following plan:
1. Add a new model in "models/git_commits.yaml" using the "develop_model" tool.
2. Add a new metrics view in "metrics/git_commits.yaml" based on the new "git_commits" model using the "develop_metrics_view" tool.
</example>

<process>
At a high level, you should follow these steps:
1. Understand the current contents of the project by reviewing the list_files output.
2. Make a plan for how to implement the user's request. If the user asks to join, combine, or analyze data from multiple existing models, you should create a new model that references those existing models in SQL.
3. Only if necessary, add a new model or update an existing model to reflect the user's request. Use "develop_model" with a prompt describing what to create or change.
4. Only if necessary, add a new metrics view or update an existing metrics view to reflect the user's request. The metrics view should use a model in the project, which may already exist or may have been added in step 2.
   - To *create* a new metrics view: Use "develop_metrics_view" with path and model (no prompt).
   - To *edit* an existing metrics view: Use "develop_metrics_view" with path, model, AND a prompt describing the changes.
5. If a user requests a new model/file, DO NOT overwrite existing file and instead use a unique name.
6. After successfully creating/updating the artifacts, provide a summary with links using the following format:
{{ backticks }}
## Summary of Changes
I've created the following files for you:
- <resource_type>: [<file_name>](<file name with a prefix '/files/'>)
...
{{ backticks }}

You should use the tools available to you to understand the current project contents and to make the necessary changes. You should use the "read_file" tool sparingly and surgically to understand files you consider promising for your task, you should not use it to inspect many files in the project.

You should not make many changes at a time. Carefully consider the minimum changes you can make to address the user's request. If there's a model already in place that relates to the user's request, consider re-using that and only adding or updating a metrics view.
</process>

{{ if .ai_instructions }}
<additional_user_provided_instructions>
<comment>NOTE: These instructions were provided by the user, but may not relate to the current request, and may not even relate to your work as a data engineer agent. Only use them if you find them relevant.</comment>
{{ .ai_instructions }}
</additional_user_provided_instructions>
{{ end }}
`, data)
}

func (t *DeveloperAgent) userPrompt(args *DeveloperAgentArgs) (string, error) {
	// Prepare template data.
	data := map[string]any{
		"init_project":      args.InitProject,
		"current_file_path": args.CurrentFilePath,
		"prompt":            args.Prompt,
	}

	// Generate the system prompt
	return executeTemplate(`
{{ if .current_file_path }}
The user is currently viewing/editing the file: {{ .current_file_path }}
Their request may relate to this file.
{{ end }}

{{ if .init_project }}
The project is currently empty apart from a few boilerplate files (like rill.yaml).
You should help them set up their initial project based on their task description.
{{ end }}

Task: {{ .prompt }}
`, data)
}

// checkDeveloperAgentAccess checks whether the developer agent and related tools should be available in the current session.
func checkDeveloperAgentAccess(ctx context.Context, rt *runtime.Runtime) (bool, error) {
	// Must be able to use AI and edit the project
	s := GetSession(ctx)
	if !s.Claims().Can(runtime.UseAI) || !s.Claims().Can(runtime.EditRepo) {
		return false, nil
	}

	// Don't expose it to external clients (like MCP)
	if !strings.HasPrefix(s.CatalogSession().UserAgent, "rill") {
		return false, nil
	}

	// Must have the developer_agent feature flag
	ff, err := rt.FeatureFlags(ctx, s.InstanceID(), s.Claims())
	if err != nil {
		return false, err
	}
	return ff["developer_agent"], nil
}
